package com.baomidou.mybatisplus.ext.processor;

import com.baomidou.mybatisplus.ext.util.AstUtil;
import com.baomidou.mybatisplus.ext.anno.Wrapperable;
import com.google.auto.service.AutoService;
import com.sun.source.tree.Tree;
import com.sun.source.util.Trees;
import com.sun.tools.javac.api.JavacTrees;
import com.sun.tools.javac.code.Flags;
import com.sun.tools.javac.model.JavacElements;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.tree.TreeTranslator;
import com.sun.tools.javac.util.*;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import java.util.Set;
import java.util.TreeSet;

@SupportedAnnotationTypes("com.baomidou.mybatisplus.ext.anno.Wrapperable")
@SupportedSourceVersion(SourceVersion.RELEASE_8)
@AutoService(Processor.class)
public class WrapperAnnotationProcessor extends AbstractProcessor {
    private Messager messager;

    private Context context;

    private JavacElements elementUtils;

    private TreeMaker treeMaker;

    private JavacTrees javacTrees;

    private Filer mFiler;

    private Set<String> stringSet = new TreeSet<>();

    private Names names;

    private Trees trees;

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        trees = Trees.instance(processingEnv);
        messager = processingEnv.getMessager();
        context = ((JavacProcessingEnvironment) processingEnv).getContext();
        elementUtils = (JavacElements) processingEnv.getElementUtils();
        treeMaker = TreeMaker.instance(context);
        javacTrees = JavacTrees.instance(processingEnv);
        mFiler = processingEnv.getFiler();
        super.init(processingEnv);
        names = Names.instance(context);
    }


    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        //获取注解标识的类集合
        Set<? extends Element> set = roundEnv.getElementsAnnotatedWith(Wrapperable.class);
        for (Element element : set) {
            //获取语法树的根节点
            JCTree tree = javacTrees.getTree(element);
            //添加访问事件
            tree.accept(new TreeTranslator() {
                @Override
                public void visitClassDef(JCTree.JCClassDecl jcClassDecl) {
                    List<JCTree.JCExpression> constructorParams = List.nil();
                    //定义一个变量叫做  wrapperVars 并且new一个ArrayList
                    jcClassDecl.defs = jcClassDecl.defs.prepend(AstUtil.addField(treeMaker, elementUtils, names,
                            "com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper",
                            "wrapper", constructorParams.append(treeMaker.Ident(names.fromString("this")))));
                    jcClassDecl.defs = jcClassDecl.defs.prepend(genAsWrapperMethod(jcClassDecl));
                    //
                    jcClassDecl.defs.stream()
                            //开始过滤: 只对变量进行处理
                            .filter(it -> it.getKind().equals(Tree.Kind.VARIABLE))
                            //强制类型转换
                            .map(it -> (JCTree.JCVariableDecl) it)
                            //处理: 生成Getter方法和Setter方法
                            .forEach(it -> {
                                        if (isHandle(it, jcClassDecl)) {
                                            //拼接通用方法
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genEQMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genNEQMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genGEMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genGTMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genLEMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genLTMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genOrderByAscMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genOrderByDescMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genIsNullMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genNotNullMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genInMethod(it, jcClassDecl));
                                            jcClassDecl.defs = jcClassDecl.defs.prepend(genNotInMethod(it, jcClassDecl));
                                            //字符串特有的like 方法拼接
                                            if (checkIsCharset(it, jcClassDecl)) {
                                                jcClassDecl.defs = jcClassDecl.defs.prepend(genLikeMethod(it, jcClassDecl));
                                                jcClassDecl.defs = jcClassDecl.defs.prepend(genLikeRightMethod(it, jcClassDecl));
                                                jcClassDecl.defs = jcClassDecl.defs.prepend(genLikeLeftMethod(it, jcClassDecl));
                                                jcClassDecl.defs = jcClassDecl.defs.prepend(genNotLikeMethod(it, jcClassDecl));
                                            }
                                        }

                                    }
                            );
                    super.visitClassDef(jcClassDecl);
                }
            });
        }
        return true;
    }


    /**
     * 构造集合 参数的方法 返回this
     *
     * @param jcVariableDecl 变量
     * @param jcClassDecl    类
     * @param methodNameStr  方法名
     * @param operator       操作符
     * @return
     */
    private JCTree.JCMethodDecl genParamCollectionMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl,
                                                         String methodNameStr, String operator) {
        // 生成 调用方法的参数列表
        List<JCTree.JCExpression> execFunparamsList = List.nil();
        // 参数添加this.wrapper
        execFunparamsList = execFunparamsList.append(treeMaker.Select(
                treeMaker.Ident(names.fromString("this"))
                , names.fromString("wrapper")
        ));

        // 操作符
        execFunparamsList = execFunparamsList.append(treeMaker.Literal(operator));

        // po
        execFunparamsList = execFunparamsList.append(treeMaker.Ident(names.fromString("this")));

        //属性名
        execFunparamsList = execFunparamsList.append(treeMaker.Literal(jcVariableDecl.name.toString()));

        //变量
        execFunparamsList = execFunparamsList.append(treeMaker.Ident(names.fromString("collection")));

        //方法参数: 方法的形式参数 Collection<?> collection
        JCTree.JCVariableDecl param = treeMaker.VarDef(
                treeMaker.Modifiers(Flags.PARAMETER, List.nil()),  //访问变量标识符
                names.fromString("collection"),        //变量名
                AstUtil.memberAccess(treeMaker, elementUtils, "java.util.Collection"),     //变量类型
                null        //变量初始值
        );

        List<JCTree.JCVariableDecl> genMethodParamsList = List.of(param);

        //生成return语句: return this;
        JCTree.JCReturn returnStatement = treeMaker.Return(
                treeMaker.Ident(names.fromString("this"))
        );
        String execMethodAllName = "com.baomidou.mybatisplus.ext.util.Funtions.addWrapper";
        return genMethod(
                 execFunparamsList
                , genMethodParamsList,
                new String[]{"com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper", "java.lang.String",
                        "java.lang.Object", "java.lang.String", "java.lang.Object"},
                treeMaker.Ident(jcClassDecl.getSimpleName()), returnStatement, execMethodAllName,names.fromString(jcVariableDecl.getName() + methodNameStr));
    }


    /**
     * 构造单个参数或者没有参数的方法
     *
     * @param jcVariableDecl 变量
     * @param jcClassDecl    类
     * @param isNeedParam    是否需要参数
     * @param methodNameStr  方法名
     * @param operator       操作符
     * @return
     */
    private JCTree.JCMethodDecl genMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl,
                                          boolean isNeedParam, String methodNameStr, String operator) {
        // 生成 调用方法的参数列表
        List<JCTree.JCExpression> execFunparamsList = List.nil();
        // 参数添加this.wrapperVars
        execFunparamsList = execFunparamsList.append(treeMaker.Select(
                treeMaker.Ident(names.fromString("this"))
                , names.fromString("wrapper")
        ));
        // 操作符
        execFunparamsList = execFunparamsList.append(treeMaker.Literal(operator));

        // po
        execFunparamsList = execFunparamsList.append(treeMaker.Ident(names.fromString("this")));

        //属性名
        execFunparamsList = execFunparamsList.append(treeMaker.Literal(jcVariableDecl.name.toString()));

        //变量
        execFunparamsList = execFunparamsList.append(isNeedParam ? treeMaker.Ident(jcVariableDecl.getName()) : treeMaker.Literal(true));

        //参数值列表
        List<JCTree.JCVariableDecl> genMethodParamsList = List.nil();
        if (isNeedParam) {
            //方法参数: 方法的形式参数
            JCTree.JCVariableDecl param = treeMaker.VarDef(
                    treeMaker.Modifiers(Flags.PARAMETER, List.nil()),  //访问变量标识符
                    jcVariableDecl.name,        //变量名
                    jcVariableDecl.vartype,     //变量类型
                    null        //变量初始值
            );

            genMethodParamsList = List.of(param);
        }

        //生成return语句: return this;
        JCTree.JCReturn returnStatement = treeMaker.Return(
                treeMaker.Ident(names.fromString("this"))
        );
        String execMethodAllName = "com.baomidou.mybatisplus.ext.util.Funtions.addWrapper";
        return genMethod(
                 execFunparamsList
                , genMethodParamsList,
                new String[]{"com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper",
                        "java.lang.String", "java.lang.Object",
                        "java.lang.String","java.lang.Object"},
                treeMaker.Ident(jcClassDecl.getSimpleName()), returnStatement, execMethodAllName,names.fromString(jcVariableDecl.getName() + methodNameStr));
    }

    /**
     * 构造方法 方法最多只有2行代码，第一行是执行某方法传一些参数，第二个是return xx
     *
     * @param execFunparamsList    调用其他方法
     * @param genMethodParamsList
     * @param execMethodParamTypes
     * @param returnMethodType
     * @param returnStatement
     * @param execMethodAllName
     * @return
     */
    private JCTree.JCMethodDecl genMethod( List<JCTree.JCExpression> execFunparamsList
            , List<JCTree.JCVariableDecl> genMethodParamsList, String[] execMethodParamTypes,
                                          JCTree.JCExpression returnMethodType,
                                          JCTree.JCReturn returnStatement, String execMethodAllName,Name methodName) {

        //语句列表  先执行语句在返回
        ListBuffer<JCTree.JCStatement> statments = new ListBuffer<JCTree.JCStatement>();
        if (execFunparamsList != null) {
            JCTree.JCMethodInvocation callMethod = AstUtil.callMethod(treeMaker, elementUtils, execMethodAllName,
                    execMethodParamTypes, execFunparamsList);
            statments.append(treeMaker.Exec(callMethod));
        }
        statments.append(returnStatement);

        //方法访问修饰符: public
        JCTree.JCModifiers modifiers = treeMaker.Modifiers(Flags.PUBLIC);


        //生成方法体
        JCTree.JCBlock body = treeMaker.Block(0, statments.toList());
        //泛型参数列表
        List<JCTree.JCTypeParameter> methodGenericParamList = List.nil();

        //异常抛出列表
        List<JCTree.JCExpression> thrownCauseList = List.nil();
        return treeMaker.MethodDef(
                modifiers,
                methodName,
                returnMethodType,
                methodGenericParamList,
                genMethodParamsList,
                thrownCauseList,
                body,
                null
        );
    }

    /**
     * 判断是否是字符串 字符串可以like
     *
     * @param jcVariableDecl 字段
     * @param jcClassDecl    类
     * @return
     */
    public boolean checkIsCharset(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {

        String typeName = jcVariableDecl.vartype.type.toString();
        try {
            Class clazz = Class.forName(typeName);
            return CharSequence.class.isAssignableFrom(clazz);
        } catch (ClassNotFoundException e) {
            return false;
        }
    }

    /**
     * 判断是否需要处理此字段
     *
     * @param jcVariableDecl 字段
     * @param jcClassDecl    类
     * @return
     */
    public boolean isHandle(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        if (jcVariableDecl.vartype == null || jcVariableDecl.vartype.type == null) {
            return false;
        }
        String typeName = jcVariableDecl.vartype.type.toString();
        //不是Java开头的不处理
        if (!typeName.startsWith("java.")) {
            return false;
        }
        try {
            Class clazz = Class.forName(typeName);
            return Comparable.class.isAssignableFrom(clazz);
        } catch (ClassNotFoundException e) {
            return false;
        }
    }


    /**
     * 生成eq方法
     *
     * @param jcVariableDecl 属性
     * @param jcClassDecl    类
     * @return
     */
    private JCTree.JCMethodDecl genEQMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "EQ", "eq");
    }

    /**
     * 生成neq方法
     *
     * @param jcVariableDecl 属性
     * @param jcClassDecl    类
     * @return
     */
    private JCTree.JCMethodDecl genNEQMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "NEQ", "neq");
    }

    /**
     * 生成like方法
     *
     * @param jcVariableDecl 属性
     * @param jcClassDecl    类
     * @return
     */
    private JCTree.JCMethodDecl genLikeMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "Like", "like");
    }

    /**
     * 生成LikeRight方法
     *
     * @param jcVariableDecl 属性
     * @param jcClassDecl    类
     * @return
     */
    private JCTree.JCMethodDecl genLikeRightMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "LikeRight", "likeRight");
    }

    /**
     * 生成LikeLeft方法
     *
     * @param jcVariableDecl 属性
     * @param jcClassDecl    类
     * @return
     */
    private JCTree.JCMethodDecl genLikeLeftMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "LikeLeft", "likeLeft");
    }

    /**
     * 生成NotLike方法
     *
     * @param jcVariableDecl 属性
     * @param jcClassDecl    类
     * @return
     */
    private JCTree.JCMethodDecl genNotLikeMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "NotLike", "notLike");
    }


    private JCTree.JCMethodDecl genGEMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "GE", "ge");
    }


    private JCTree.JCMethodDecl genGTMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "GT", "gt");
    }

    private JCTree.JCMethodDecl genLEMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "LE", "le");
    }

    private JCTree.JCMethodDecl genLTMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                true, "LT", "lt");
    }

    private JCTree.JCMethodDecl genOrderByAscMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                false, "OrderByAsc", "orderByAsc");
    }

    private JCTree.JCMethodDecl genOrderByDescMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                false, "OrderByDesc", "orderByDesc");
    }

    private JCTree.JCMethodDecl genIsNullMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                false, "IsNull", "isNull");
    }

    private JCTree.JCMethodDecl genNotNullMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genMethod(jcVariableDecl, jcClassDecl,
                false, "NotNull", "notNull");
    }

    private JCTree.JCMethodDecl genInMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genParamCollectionMethod(jcVariableDecl, jcClassDecl, "In", "in");
    }

    private JCTree.JCMethodDecl genNotInMethod(JCTree.JCVariableDecl jcVariableDecl, JCTree.JCClassDecl jcClassDecl) {
        return genParamCollectionMethod(jcVariableDecl, jcClassDecl, "NotIn", "notIn");
    }

    /**
     * 生成asWrapper方法
     * @param jcClassDecl 类
     * @return
     */
    private JCTree.JCMethodDecl genAsWrapperMethod(JCTree.JCClassDecl jcClassDecl) {
        //生成return语句: return this;
        JCTree.JCReturn returnStatement = treeMaker.Return(
                treeMaker.Select(
                        treeMaker.Ident(names.fromString("this"))
                        , names.fromString("wrapper")
                )
        );
        return genMethod(null
                , List.nil(),
                null,
                AstUtil.memberAccess(treeMaker,elementUtils,"com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper"),
                returnStatement, null, names.fromString("asWrapper"));
    }


}
